{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "27a1b7fe",
   "metadata": {},
   "source": [
    "## Mnist分类任务\n",
    "\n",
    "- 网络基本构建与训练方法，常用函数解析\n",
    "- torch.nn.functional模块\n",
    "- nn.Module模块"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50594e27",
   "metadata": {},
   "source": [
    "## 读取Mnist数据集\n",
    "\n",
    "- 会自动进行下载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "105d35ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5240c6f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import requests\n",
    "\n",
    "# 下载数据集\n",
    "DATA_PATH = Path('data')\n",
    "PATH = DATA_PATH / 'mnist'\n",
    "\n",
    "PATH.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "URL = 'http://deeplearning.net/data/mnist'\n",
    "FILENAME = 'mnist.pkl.gz'\n",
    "\n",
    "if not (PATH / FILENAME).exists():\n",
    "    content = requests.get(URL + FILENAME).content\n",
    "    (PATH / FILENAME).open('wb').write(content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "452f8862",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import gzip\n",
    "\n",
    "with gzip.open((PATH / FILENAME).as_posix(), 'rb') as f:\n",
    "    ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding='latin-1')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7c881f5",
   "metadata": {},
   "source": [
    "- 784是mnist数据集每个样本的像素点个数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10b8114c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot\n",
    "import numpy as np\n",
    "\n",
    "pyplot.inshow(x_train[0].reshape((28, 28)), cmap='gray')\n",
    "print(x_train.shape)\n",
    "\n",
    "# (50000, 784)\n",
    "# 50000个数据样本，28x28x1  28 x 28 矩阵 1个颜色通道"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "264206e1",
   "metadata": {},
   "source": [
    "- 注意：数据需转换成tensor才能参与后续建模训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a80d144a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "x_train, y_train, x_valid, y_valid = map(\n",
    "    torch.tensor, (x_train, y_train, x_valid, y_valid)\n",
    ")\n",
    "\n",
    "n, c = x_train.shape\n",
    "x_train, x_train.shape, y_train.min(), y_train.max()\n",
    "\n",
    "print(x_train, y_train)\n",
    "print(x_train.shape)\n",
    "print(y_train.min(), y_train.max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1cebb7a",
   "metadata": {},
   "source": [
    "#### torch.nn.functional 很多层和函数在这里都会见到\n",
    "\n",
    "torch.nn.functional 中有很多功能，后续会常用到。什么时候使用nn.Module，什么时候使用nn.functional呢？一般情况下，如果模型有可学习的参数，最好用nn.Module，其他情况nn.functional相对更简单一些"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "442033e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "loss_func = F.cross_entropy\n",
    "\n",
    "def model(xb):\n",
    "    return xb.mm(weights) + bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7d21fe64",
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 64\n",
    "# a mini-batch from x\n",
    "xb = x_train[0: bs]\n",
    "yb = y_train[0: bs]\n",
    "weights = torch.randn([784, 10], dtype=torch.float, requires_grad=True)\n",
    "bs = 64\n",
    "bias = torch.zeros(10, requires_grad=True)\n",
    "\n",
    "# 实际损失值\n",
    "print(loss_func(model(xb), yb))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f28b4335",
   "metadata": {},
   "source": [
    "#### 创建一个model来更简化代码\n",
    "\n",
    "- 必须基础nn.Module 且在其构造函数中需调用nn.Module的构造函数\n",
    "- 无需写反向传播函数，nn.Module能够利用autograd自动实现反向传播\n",
    "- Module中的可学习参数可以通过named_parameters() 或者parameters() 返回迭代器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23fa37ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "\n",
    "class Mnist_NN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.hidden1 = nn.Linear(784, 128)\n",
    "        self.hidden2 = nn.Linear(128, 256)\n",
    "        self.out = nn.Linear(256, 10)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.hidden1(x))\n",
    "        x = F.relu(self.hidden2(x))\n",
    "        x = self.out(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f5ba052",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Mnist_NN()\n",
    "print(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50faca6b",
   "metadata": {},
   "source": [
    "- 可以打印我们定义好名字里的权重的偏置项"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7ff4387",
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, parameter in net.named_parameters():\n",
    "    print(name, parameter, parameter.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c60e096",
   "metadata": {},
   "source": [
    "## 使用 TensorDataset和DataLoader来简化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54997fea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "train_ds = TensorDataset(x_train, y_train)\n",
    "train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)\n",
    "\n",
    "valid_ds = TensorDataset(x_valid, y_valid)\n",
    "valid_dl = TataLoader(valid_ds, batch_size=bs * 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f442471e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(train_ds, valid_ds, bs):\n",
    "    return (\n",
    "        DataLoader(train_ds, batch_size=bs, shuffle=True),\n",
    "        DataLoader(train_ds, batch_size=bs * 2)\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f834463c",
   "metadata": {},
   "source": [
    "- 一般在训练模型时加上model.train()，这样会正常使用Batch Normalization 和Dropout\n",
    "- 测试的时候一般选择model.eval()，这样就不会使用Batch Normalization 和Dropout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b8a1109",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def fit(steps, model, loss_func, opt, train_dl, valid_dl):\n",
    "    for step in range(steps):\n",
    "        model.train()\n",
    "        for xb, yb in train_dl:\n",
    "            loss_batch(model, loss_func, xb, yb, opt)\n",
    "        \n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            losses, nums = zip(\n",
    "                *[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]\n",
    "            )\n",
    "        # 验证集损失\n",
    "        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)\n",
    "        print('当前step：' + str(step), '验证集损失：' + str(val_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44444cfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import optim\n",
    "\n",
    "def get_model():\n",
    "    model = Mnist_NN()\n",
    "    return model, optim.SGD(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae96bf0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_batch(model, loss_func, xb, yb, opt=None):\n",
    "    # 计算当前损失，model(xb)为预测结果，yb为真实值\n",
    "    loss = loss_func(model(xb), yb)\n",
    "    \n",
    "    if opt is not None:\n",
    "        # 求梯度\n",
    "        loss.backward()\n",
    "        # 更新\n",
    "        opt.step()\n",
    "        # 置零，清空\n",
    "        opt.zero_grad()\n",
    "        \n",
    "    return loss.item(), len(xb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80ffcb08",
   "metadata": {},
   "source": [
    "三行搞定！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6017dad",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dl, valid_dl = get_data(train_ds, valid_ds, bs)\n",
    "model, opt = get_model()\n",
    "fit(25, model, loss_func, opt, train_dl, valid_dl)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
